{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "[YOLOv8 of roboflow](https://colab.research.google.com/github/QIN2DIM/hcaptcha-challenger/blob/main/automation/roboflow_yolov8.ipynb)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tdSMcABDNKW-",
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "outputId": "0380de4e-243c-4bb6-9e85-14ff6d4e9a57"
   },
   "outputs": [],
   "source": [
    "!pip install -U onnx onnxruntime-gpu onnxruntime\n",
    "!pip install ultralytics==8.0.202"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "Paste roboflow snippet into here ↓"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# !pip install roboflow\n",
    "#\n",
    "# from roboflow import Roboflow\n",
    "# rf = Roboflow(api_key=\"APIKEY\")\n",
    "# project = rf.workspace(\"WORKSPACE\").project(\"PROJECT\")\n",
    "# dataset = project.version(1).download(\"yolov8\")"
   ],
   "metadata": {
    "id": "oaI2rzp4pppS",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "outputId": "10bdc289-e060-4eb0-974a-29d929f4cfae"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "Train"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "from pathlib import Path\n",
    "\n",
    "from ultralytics import YOLO\n",
    "\n",
    "HOME = os.getcwd()\n",
    "%cd {HOME}\n",
    "\n",
    "challenge_dir_name = \"animal-head-2310-1\"\n",
    "output_model_name = \"head_of_the_animal_2311_yolov8s.onnx\"\n",
    "\n",
    "this_dir = Path(os.path.abspath(\"__file__\")).parent\n",
    "yolo_data_dir = this_dir.joinpath(challenge_dir_name)\n",
    "yolo_inspect_path = yolo_data_dir.joinpath(\"data.yaml\")\n",
    "\n",
    "model = YOLO(\"yolov8s.pt\")  # detection\n",
    "# model = YOLO(\"yolov8s-seg.pt\") # segment\n",
    "model.train(data=str(yolo_inspect_path), epochs=300, imgsz=640)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Export YOLOv8 to ONNX"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# detection\n",
    "yolo_best_pt_path = this_dir.joinpath(\"runs/detect/train/weights/best.pt\")\n",
    "yolo_onnx_path = this_dir.joinpath(\"runs/detect/train/weights/best.onnx\")\n",
    "\n",
    "# segment\n",
    "# yolo_best_pt_path = this_dir.joinpath(\"runs/segment/train/weights/best.pt\")\n",
    "# yolo_onnx_path = this_dir.joinpath(\"runs/segment/train/weights/best.onnx\")\n",
    "\n",
    "model_path = this_dir.joinpath(output_model_name)\n",
    "\n",
    "# export to onnx\n",
    "YOLO(str(yolo_best_pt_path)).export(format=\"onnx\", simplify=True, int8=True)\n",
    "\n",
    "# copy to current work dir\n",
    "shutil.copyfile(yolo_onnx_path, model_path)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "Deployment model to GitHub"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import locale\n",
    "\n",
    "locale.getpreferredencoding = lambda d=True: \"UTF-8\"\n",
    "\n",
    "!pip install PyGithub\n",
    "\n",
    "from github import Auth, Github, GithubException\n",
    "\n",
    "auth = Auth.Token(os.getenv(\"GITHUB_TOKEN\", \"\"))\n",
    "repo = Github(auth=auth).get_repo(\"QIN2DIM/hcaptcha-challenger\")\n",
    "modelhub_title = \"ONNX ModelHub\"\n",
    "\n",
    "\n",
    "def quick_development(pending_onnx_path: Path):\n",
    "    if not pending_onnx_path.exists() or not pending_onnx_path.name.endswith(\".onnx\"):\n",
    "        return\n",
    "\n",
    "    for release in repo.get_releases():\n",
    "        if release.title != modelhub_title:\n",
    "            continue\n",
    "        try:\n",
    "            res = release.upload_asset(path=str(pending_onnx_path))\n",
    "        except GithubException as err:\n",
    "            if err.status == 422:\n",
    "                print(\n",
    "                    \"The model file already exists, please manually replace the file with the same\"\n",
    "                    f\" name - url={repo.releases_url}\"\n",
    "                )\n",
    "        except Exception as err:\n",
    "            print(err)\n",
    "        else:\n",
    "            print(\n",
    "                f\"Model file uploaded successfully - name={res.name} url={res.browser_download_url}\"\n",
    "            )\n",
    "\n",
    "\n",
    "quick_development(pending_onnx_path=model_path)"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "name": "python3",
   "language": "python",
   "display_name": "Python 3 (ipykernel)"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
